import torch
import torchvision
import numpy as np
import d2lzh_pytorch as d2l

# 读取数据
batch_size = 256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)

num_inputs = 784
num_outputs = 10

# 一共28 * 28 = 784个像素, 10个类别
# 784 * 10
W = torch.tensor(np.random.normal(0, 0.01, (num_inputs, num_outputs)), dtype=torch.float)
# 1 * 10
b = torch.zeros(num_outputs, dtype=torch.float)

W.requires_grad_(requires_grad=True)
b.requires_grad_(requires_grad=True)

# X的行数代表样本数, 列数代表输出个数(类别数)
def softmax(X):
    X_exp = X.exp()
    # dim=1代表对 行 这个维度进行操作, keepdim = True代表操作完后保留行和列的维度
    partition = X_exp.sum(dim=1, keepdim=True)
    return X_exp / partition  # 这里应用了广播机制


# 定义模型
def net(X):
    # 将每张原始图像改成长度为num_inputs的向量
    return softmax(torch.mm(X.view((-1, num_inputs)), W) + b)


# 定义损失函数
"""
# gather函数的作用: 收集输入的特定维度指定位置的数值
# y_hat是2个样本在3个类别的预测概率，变量y是这2个样本的标签类别
# 2 * 3
y_hat = torch.tensor([[0.1, 0.3, 0.6], [0.3, 0.2, 0.5]])
# 1 * 2
y = torch.LongTensor([0, 2])
print(y_hat.gather(1, y.view(-1, 1)))
"""
# 交叉熵损失
def cross_entropy(y_hat, y):
    return - torch.log(y_hat.gather(1, y.view(-1, 1)))

# 计算分类准确率
# 给定一个类别的预测概率分布y_hat，我们把预测概率最大的类别作为输出类别。如果它与真实类别y一致，说明这次预测是正确的。分类准确率即正确预测数量与总预测数量之比
def accuracy(y_hat, y):
    return (y_hat.argmax(dim=1) == y).float().mean().item()

# print(d2l.evaluate_accuracy(test_iter, net))

# 训练模型
num_epochs, lr = 5, 0.1

d2l.train_ch3(net, train_iter, test_iter, cross_entropy, num_epochs, batch_size, [W, b], lr)


# 预测
X, y = iter(test_iter).next()

true_labels = d2l.get_fashion_mnist_labels(y.numpy())
pred_labels = d2l.get_fashion_mnist_labels(net(X).argmax(dim=1).numpy())
titles = [true + '\n' + pred for true, pred in zip(true_labels, pred_labels)]

d2l.show_fashion_mnist(X[0:9], titles[0:9])
